// Copyright Epic Games, Inc. All Rights Reserved.

#pragma once

#include <memory.h>
#include <zenbase/refcount.h>
#include <zencore/memory.h>
#include <atomic>
#include "zencore.h"

#include <filesystem>

namespace zen {

struct IoHash;
struct IoBufferExtendedCore;

enum class ZenContentType : uint8_t
{
	kBinary				= 0,  // Note that since this is zero, this will be the default value in IoBuffer
	kText				= 1,
	kJSON				= 2,
	kCbObject			= 3,
	kCbPackage			= 4,
	kYAML				= 5,
	kCbPackageOffer		= 6,
	kCompressedBinary	= 7,
	kUnknownContentType = 8,
	kHTML				= 9,
	kJavaScript			= 10,
	kCSS				= 11,
	kPNG				= 12,
	kIcon				= 13,
	kXML				= 14,
	kCOUNT
};

inline std::string_view
ToString(ZenContentType ContentType)
{
	using namespace std::literals;

	switch (ContentType)
	{
		default:
		case ZenContentType::kUnknownContentType:
			return "unknown"sv;
		case ZenContentType::kBinary:
			return "binary"sv;
		case ZenContentType::kText:
			return "text"sv;
		case ZenContentType::kJSON:
			return "json"sv;
		case ZenContentType::kCbObject:
			return "cb-object"sv;
		case ZenContentType::kCbPackage:
			return "cb-package"sv;
		case ZenContentType::kCbPackageOffer:
			return "cb-package-offer"sv;
		case ZenContentType::kCompressedBinary:
			return "compressed-binary"sv;
		case ZenContentType::kYAML:
			return "yaml"sv;
		case ZenContentType::kHTML:
			return "html"sv;
		case ZenContentType::kJavaScript:
			return "javascript"sv;
		case ZenContentType::kCSS:
			return "css"sv;
		case ZenContentType::kPNG:
			return "png"sv;
		case ZenContentType::kIcon:
			return "icon"sv;
		case ZenContentType::kXML:
			return "xml"sv;
	}
}

struct IoBufferFileReference
{
	void*	 FileHandle;
	uint64_t FileChunkOffset;
	uint64_t FileChunkSize;
};

struct IoBufferCore
{
public:
	inline IoBufferCore() : m_Flags(kIsNull) {}
	inline IoBufferCore(const void* DataPtr, size_t SizeBytes) : m_DataPtr(DataPtr), m_DataBytes(SizeBytes) {}
	inline IoBufferCore(const IoBufferCore* Outer, const void* DataPtr, size_t SizeBytes)
	: m_DataPtr(DataPtr)
	, m_DataBytes(SizeBytes)
	, m_OuterCore(Outer)
	{
	}

	ZENCORE_API explicit IoBufferCore(size_t SizeBytes);
	ZENCORE_API IoBufferCore(size_t SizeBytes, size_t Alignment);
	ZENCORE_API ~IoBufferCore();

	// Reference counting

	inline uint32_t AddRef() const { return AtomicIncrement(const_cast<IoBufferCore*>(this)->m_RefCount); }
	inline uint32_t Release() const
	{
		const uint32_t NewRefCount = AtomicDecrement(const_cast<IoBufferCore*>(this)->m_RefCount);
		if (NewRefCount == 0)
		{
			DeleteThis();
		}
		return NewRefCount;
	}

	// Copying reference counted objects doesn't make a lot of sense generally, so let's prevent it

	IoBufferCore(const IoBufferCore&) = delete;
	IoBufferCore(IoBufferCore&&)	  = delete;
	IoBufferCore& operator=(const IoBufferCore&) = delete;
	IoBufferCore& operator=(IoBufferCore&&) = delete;

	//

	ZENCORE_API void Materialize() const;
	ZENCORE_API void DeleteThis() const;
	ZENCORE_API void MakeOwned(bool Immutable = true);

	inline void EnsureDataValid() const
	{
		const uint32_t LocalFlags = m_Flags.load(std::memory_order_acquire);
		if ((LocalFlags & kIsExtended) && !(LocalFlags & kIsMaterialized))
		{
			Materialize();
		}
	}

	inline bool IsOwnedByThis() const { return !!(m_Flags.load(std::memory_order_relaxed) & kIsOwnedByThis); }

	inline void SetIsOwnedByThis(bool NewState)
	{
		if (NewState)
		{
			m_Flags.fetch_or(kIsOwnedByThis, std::memory_order_relaxed);
		}
		else
		{
			m_Flags.fetch_and(~kIsOwnedByThis, std::memory_order_relaxed);
		}
	}

	inline bool IsOwned() const
	{
		if (IsOwnedByThis())
		{
			return true;
		}
		return m_OuterCore && m_OuterCore->IsOwned();
	}

	inline bool IsImmutable() const { return (m_Flags.load(std::memory_order_relaxed) & kIsMutable) == 0; }
	inline bool IsWholeFile() const { return (m_Flags.load(std::memory_order_relaxed) & kIsWholeFile) != 0; }
	inline bool IsNull() const { return (m_Flags.load(std::memory_order_relaxed) & kIsNull) != 0; }

	inline IoBufferExtendedCore*	   ExtendedCore();
	inline const IoBufferExtendedCore* ExtendedCore() const;

	ZENCORE_API void* MutableDataPointer() const;

	inline const void* DataPointer() const
	{
		EnsureDataValid();
		return m_DataPtr;
	}

	inline size_t DataBytes() const { return m_DataBytes; }

	inline void Set(const void* Ptr, size_t Sz)
	{
		m_DataPtr	= Ptr;
		m_DataBytes = Sz;
	}

	inline void SetIsImmutable(bool NewState)
	{
		if (!NewState)
		{
			m_Flags.fetch_or(kIsMutable, std::memory_order_relaxed);
		}
		else
		{
			m_Flags.fetch_and(~kIsMutable, std::memory_order_relaxed);
		}
	}

	inline void SetIsWholeFile(bool NewState)
	{
		if (NewState)
		{
			m_Flags.fetch_or(kIsWholeFile, std::memory_order_relaxed);
		}
		else
		{
			m_Flags.fetch_and(~kIsWholeFile, std::memory_order_relaxed);
		}
	}

	inline void SetContentType(ZenContentType ContentType)
	{
		ZEN_ASSERT_SLOW((uint32_t(ContentType) & kContentTypeMask) == uint32_t(ContentType));
		uint32_t OldValue = m_Flags.load(std::memory_order_relaxed);
		uint32_t NewValue;
		do
		{
			NewValue = (OldValue & ~(kContentTypeMask << kContentTypeShift)) | (uint32_t(ContentType) << kContentTypeShift);
		} while (!m_Flags.compare_exchange_weak(OldValue, NewValue, std::memory_order_relaxed, std::memory_order_relaxed));
	}

	inline ZenContentType GetContentType() const
	{
		return ZenContentType((m_Flags.load(std::memory_order_relaxed) >> kContentTypeShift) & kContentTypeMask);
	}

	inline uint32_t GetRefCount() const { return m_RefCount; }

protected:
	uint32_t					  m_RefCount = 0;
	mutable std::atomic<uint32_t> m_Flags{0};
	mutable const void*			  m_DataPtr	  = nullptr;
	size_t						  m_DataBytes = 0;
	RefPtr<const IoBufferCore>	  m_OuterCore;

	enum
	{
		kContentTypeShift = 24,
		kContentTypeMask  = 0xf
	};

	static_assert((uint32_t(ZenContentType::kUnknownContentType) & ~kContentTypeMask) == 0);

	enum Flags : uint32_t
	{
		kIsNull			= 1 << 0,  // This is a null IoBuffer
		kIsMutable		= 1 << 1,
		kIsExtended		= 1 << 2,  // Is actually a SharedBufferExtendedCore
		kIsMaterialized = 1 << 3,  // Data pointers are valid
		kLowLevelAlloc	= 1 << 4,  // Using direct memory allocation
		kIsWholeFile	= 1 << 5,  // References an entire file
		kIoBufferAlloc	= 1 << 6,  // Using IoBuffer allocator
		kIsOwnedByThis	= 1 << 7,

		// Note that we have some extended flags defined below
		// so not all bits are available to use here

		kContentTypeBit0 = 1 << (24 + 0),  // These constants
		kContentTypeBit1 = 1 << (24 + 1),  // are here mostly to
		kContentTypeBit2 = 1 << (24 + 2),  // indicate that these
		kContentTypeBit3 = 1 << (24 + 3),  // bits are reserved
	};

	void AllocateBuffer(size_t InSize, size_t Alignment) const;
	void FreeBuffer();
};

/**
 * An "Extended" core references a segment of a file
 */

struct IoBufferExtendedCore : public IoBufferCore
{
	IoBufferExtendedCore(void* FileHandle, uint64_t Offset, uint64_t Size, bool TransferHandleOwnership);
	IoBufferExtendedCore(const IoBufferExtendedCore* Outer, uint64_t Offset, uint64_t Size);
	~IoBufferExtendedCore();

	enum ExtendedFlags
	{
		kOwnsFile	   = 1 << 16,
		kOwnsMmap	   = 1 << 17,
		kDeleteOnClose = 1 << 18
	};

	void Materialize() const;
	bool GetFileReference(IoBufferFileReference& OutRef) const;
	void SetDeleteOnClose(bool DeleteOnClose);

private:
	void*		  m_FileHandle	  = nullptr;
	uint64_t	  m_FileOffset	  = 0;
	mutable void* m_MmapHandle	  = nullptr;
	mutable void* m_MappedPointer = nullptr;
};

inline IoBufferExtendedCore*
IoBufferCore::ExtendedCore()
{
	if (m_Flags.load(std::memory_order_relaxed) & kIsExtended)
	{
		return static_cast<IoBufferExtendedCore*>(this);
	}

	return nullptr;
}

inline const IoBufferExtendedCore*
IoBufferCore::ExtendedCore() const
{
	if (m_Flags.load(std::memory_order_relaxed) & kIsExtended)
	{
		return static_cast<const IoBufferExtendedCore*>(this);
	}

	return nullptr;
}

/**
 * I/O buffer
 *
 * This represents a reference to a payload in memory or on disk
 *
 */
class IoBuffer
{
public:
	enum ECloneTag
	{
		Clone
	};
	enum EWrapTag
	{
		Wrap
	};
	enum EFileTag
	{
		File
	};
	enum EBorrowedFileTag
	{
		BorrowedFile
	};

	inline IoBuffer() = default;
	inline IoBuffer(IoBuffer&& Rhs) noexcept
	{
		m_Core.Swap(Rhs.m_Core);
		Rhs.m_Core = NullBufferCore;
	}
	inline IoBuffer(const IoBuffer& Rhs) = default;
	inline IoBuffer& operator=(const IoBuffer& Rhs) = default;
	inline IoBuffer& operator						=(IoBuffer&& Rhs) noexcept
	{
		m_Core.Swap(Rhs.m_Core);
		Rhs.m_Core = NullBufferCore;
		return *this;
	}

	/** Create an uninitialized buffer of the given size
	 */
	ZENCORE_API explicit IoBuffer(size_t InSize);

	/** Create an uninitialized buffer of the given size with the specified alignment
	 */
	ZENCORE_API explicit IoBuffer(size_t InSize, uint64_t InAlignment);

	/** Create a buffer which references a sequence of bytes inside another buffer
	 */
	ZENCORE_API IoBuffer(const IoBuffer& OuterBuffer, size_t Offset, size_t SizeBytes = ~0ull);

	/** Create a buffer which references a range of bytes which we assume will live
	 * for the entire life time.
	 */
	inline IoBuffer(EWrapTag, const void* DataPtr, size_t SizeBytes) : m_Core(new IoBufferCore(DataPtr, SizeBytes)) {}

	inline IoBuffer(ECloneTag, const void* DataPtr, size_t SizeBytes) : m_Core(new IoBufferCore(SizeBytes))
	{
		memcpy(const_cast<void*>(m_Core->DataPointer()), DataPtr, SizeBytes);
	}

	ZENCORE_API IoBuffer(EFileTag, void* FileHandle, uint64_t ChunkFileOffset, uint64_t ChunkSize, bool IsWholeFile);
	ZENCORE_API IoBuffer(EBorrowedFileTag, void* FileHandle, uint64_t ChunkFileOffset, uint64_t ChunkSize);

	inline explicit						operator bool() const { return !m_Core->IsNull(); }
	inline								operator MemoryView() const& { return MemoryView(m_Core->DataPointer(), m_Core->DataBytes()); }
	inline void							MakeOwned() const { return m_Core->MakeOwned(); }
	[[nodiscard]] inline bool			IsOwned() const { return m_Core->IsOwned(); }
	[[nodiscard]] inline bool			IsWholeFile() const { return m_Core->IsWholeFile(); }
	[[nodiscard]] void*					MutableData() const { return m_Core->MutableDataPointer(); }
	void								MakeImmutable() { m_Core->SetIsImmutable(true); }
	[[nodiscard]] const void*			Data() const { return m_Core->DataPointer(); }
	[[nodiscard]] const void*			GetData() const { return m_Core->DataPointer(); }
	[[nodiscard]] size_t				Size() const { return m_Core->DataBytes(); }
	[[nodiscard]] size_t				GetSize() const { return m_Core->DataBytes(); }
	inline void							SetContentType(ZenContentType ContentType) { m_Core->SetContentType(ContentType); }
	[[nodiscard]] inline ZenContentType GetContentType() const { return m_Core->GetContentType(); }
	[[nodiscard]] ZENCORE_API bool		GetFileReference(IoBufferFileReference& OutRef) const;
	void								SetDeleteOnClose(bool DeleteOnClose);

	inline MemoryView		 GetView() const { return MemoryView(m_Core->DataPointer(), m_Core->DataBytes()); }
	inline MutableMemoryView GetMutableView() { return MutableMemoryView(m_Core->MutableDataPointer(), m_Core->DataBytes()); }

	template<typename T>
	[[nodiscard]] const T* Data() const
	{
		return reinterpret_cast<const T*>(m_Core->DataPointer());
	}

	template<typename T>
	[[nodiscard]] T* MutableData() const
	{
		return reinterpret_cast<T*>(m_Core->MutableDataPointer());
	}

private:
	// We have a shared "null" buffer core which we share, this is initialized static and never released which will
	// cause a memory leak at exit. This does however save millions of memory allocations for null buffers
	static RefPtr<IoBufferCore> NullBufferCore;

	RefPtr<IoBufferCore> m_Core = NullBufferCore;

	IoBuffer(IoBufferCore* Core) : m_Core(Core) {}

	friend class SharedBuffer;
	friend class IoBufferBuilder;
};

class IoBufferBuilder
{
public:
	ZENCORE_API static IoBuffer MakeFromFile(const std::filesystem::path& FileName, uint64_t Offset = 0, uint64_t Size = ~0ull);
	ZENCORE_API static IoBuffer MakeFromTemporaryFile(const std::filesystem::path& FileName);
	ZENCORE_API static IoBuffer MakeFromFileHandle(void* FileHandle, uint64_t Offset = 0, uint64_t Size = ~0ull);
	/** Make sure buffer data is memory resident, but avoid memory mapping data from files
	 */
	ZENCORE_API static IoBuffer ReadFromFileMaybe(const IoBuffer& InBuffer);
	inline static IoBuffer		MakeFromMemory(MemoryView Memory) { return IoBuffer(IoBuffer::Wrap, Memory.GetData(), Memory.GetSize()); }
	inline static IoBuffer		MakeCloneFromMemory(const void* Ptr, size_t Sz)
	{
		if (Sz)
		{
			return IoBuffer(IoBuffer::Clone, Ptr, Sz);
		}
		return {};
	}
	inline static IoBuffer MakeCloneFromMemory(MemoryView Memory) { return MakeCloneFromMemory(Memory.GetData(), Memory.GetSize()); }
};

void iobuffer_forcelink();

}  // namespace zen
